#include "types.h"
#include "riscv.h"
#include "defs.h"
#include "param.h"
#include "memlayout.h"
#include "spinlock.h"
#include "sleeplock.h"
#include "fs.h"
#include "buf.h"
#include "virtio.h"
#include "debug.h"

#define R(r) ((volatile uint32 *)(VIRTIO0 + (r)))
#define VIRTIO_CONS_BUF_SIZE 64

static struct cons {
    char i_pages[2*PGSIZE];
    struct virtq_desc *i_desc;
    struct virtq_avail *i_avail;
    struct virtq_used *i_used;
    char o_pages[2*PGSIZE];
    struct virtq_desc *o_desc;
    struct virtq_avail *o_avail;
    struct virtq_used *o_used;
    char i_free[NUM];  // is a descriptor free?
    uint16 i_used_idx; // we've looked this far in used[2..NUM].
    char o_free[NUM];
    uint16 o_used_idx;
    struct spinlock vcons_lock;
    int out_idx[NUM];
    char inbuf[NUM][VIRTIO_CONS_BUF_SIZE];
    int inlen;
} __attribute__ ((aligned (PGSIZE))) cons;

void virtio_console_input_supply ();

static int
alloc_i_desc () {
    for (int i = 0; i < NUM; i++) {
        if (cons.i_free[i]) {
            cons.i_free[i] = 0;
            return i;
        }
    }
    return -1;
}

static void
free_i_desc (int i) {
    if (i >= NUM)
        panic ("free_i_desc");
    if (cons.i_free[i])
        panic ("free_i_desc");
    cons.i_desc[i].addr = 0;
    cons.i_desc[i].len = 0;
    cons.i_desc[i].flags = 0;
    cons.i_desc[i].next = 0;
    cons.i_free[i] = 1;
    //wakeup (&cons.i_free[0]);
}

static int
alloc_o_desc () {
    for (int i = 0; i < NUM; i++) {
        if (cons.o_free[i]) {
            cons.o_free[i] = 0;
            return i;
        }
    }
    return -1;
}

static void
free_o_desc (int i) {
    if (i >= NUM)
        panic ("free_o_desc");
    if (cons.o_free[i])
        panic ("free_o_desc");
    cons.o_desc[i].addr = 0;
    cons.o_desc[i].len = 0;
    cons.o_desc[i].flags = 0;
    cons.o_desc[i].next = 0;
    cons.o_free[i] = 0;
    wakeup (&cons.o_free[0]);
}

void
virtio_console_init (void) {
    memset (&cons, 0, sizeof (cons));
    initlock (&cons.vcons_lock, "virtio_cons");
    uint32 status = 0;
    if(*R(VIRTIO_MMIO_MAGIC_VALUE) != 0x74726976 ||
        *R(VIRTIO_MMIO_VERSION) != 2 ||
        *R(VIRTIO_MMIO_DEVICE_ID) != 3 ||
        *R(VIRTIO_MMIO_VENDOR_ID) != 0x554d4551){
        panic("could not find virtio console");
    }
    status |= VIRTIO_CONFIG_S_ACKNOWLEDGE;
    *R(VIRTIO_MMIO_STATUS) = status;

    status |= VIRTIO_CONFIG_S_DRIVER;
    *R(VIRTIO_MMIO_STATUS) = status;

    // negotiate features
    uint64 features = *R(VIRTIO_MMIO_DEVICE_FEATURES);
    /* simple console device :D */
    *R(VIRTIO_MMIO_DRIVER_FEATURES) = features;

    status |= VIRTIO_CONFIG_S_FEATURES_OK;
    *R(VIRTIO_MMIO_STATUS) = status;

    status |= VIRTIO_CONFIG_S_DRIVER_OK;
    *R(VIRTIO_MMIO_STATUS) = status;
    uint32 max;
    /* initialize queue 0 (input) */
    *R(VIRTIO_MMIO_QUEUE_SEL) = 0;
    max = *R(VIRTIO_MMIO_QUEUE_NUM_MAX);
    if(max == 0)
        panic("virtio console has no queue 0");
    if(max < NUM)
        panic("virtio console max queue too short");
    *R(VIRTIO_MMIO_QUEUE_NUM) = NUM;
    memset(cons.i_pages, 0, sizeof(cons.i_pages));
    cons.i_desc = (struct virtq_desc *) cons.i_pages;
    cons.i_avail = (struct virtq_avail *) (cons.i_pages + NUM * sizeof (struct virtq_desc));
    cons.i_used = (struct virtq_used *) (cons.i_pages + PGSIZE);
    
    /* set queue address */
    *R(VIRTIO_MMIO_QUEUE_AVAIL_HIGH) = (uint32)((uint64)cons.i_avail >> 32);
    *R(VIRTIO_MMIO_QUEUE_AVAIL_LOW) = (uint32)((uint64)cons.i_avail & 0xffffffff);

    *R(VIRTIO_MMIO_QUEUE_DESC_HIGH) = (uint32)((uint64)cons.i_desc >> 32);
    *R(VIRTIO_MMIO_QUEUE_DESC_LOW) = (uint32)((uint64)cons.i_desc & 0xffffffff);

    *R(VIRTIO_MMIO_QUEUE_USED_HIGH) = (uint32)((uint64)cons.i_used >> 32);
    *R(VIRTIO_MMIO_QUEUE_USED_LOW) = (uint32)((uint64)cons.i_used & 0xffffffff);

    for (int i = 0; i < NUM; i++) {
        cons.i_free[i] = 1;
    }

    *R(VIRTIO_MMIO_QUEUE_READY) = 1;

    /* initialize queue 1 (output) */
    *R(VIRTIO_MMIO_QUEUE_SEL) = 1;
    max = *R(VIRTIO_MMIO_QUEUE_NUM_MAX);
    if(max == 0)
        panic("virtio console has no queue 1");
    if(max < NUM)
        panic("virtio console max queue too short");
    *R(VIRTIO_MMIO_QUEUE_NUM) = NUM;
    memset(cons.o_pages, 0, sizeof(cons.o_pages));
    cons.o_desc = (struct virtq_desc *) cons.o_pages;
    cons.o_avail = (struct virtq_avail *) (cons.o_pages + NUM * sizeof (struct virtq_desc));
    cons.o_used = (struct virtq_used *) (cons.o_pages + PGSIZE);
    
    /* set queue address */
    *R(VIRTIO_MMIO_QUEUE_AVAIL_HIGH) = (uint32)((uint64)cons.o_avail >> 32);
    *R(VIRTIO_MMIO_QUEUE_AVAIL_LOW) = (uint32)((uint64)cons.o_avail & 0xffffffff);

    *R(VIRTIO_MMIO_QUEUE_DESC_HIGH) = (uint32)((uint64)cons.o_desc >> 32);
    *R(VIRTIO_MMIO_QUEUE_DESC_LOW) = (uint32)((uint64)cons.o_desc & 0xffffffff);

    *R(VIRTIO_MMIO_QUEUE_USED_HIGH) = (uint32)((uint64)cons.o_used >> 32);
    *R(VIRTIO_MMIO_QUEUE_USED_LOW) = (uint32)((uint64)cons.o_used & 0xffffffff);

    for (int i = 0; i < NUM; i++) {
        cons.o_free[i] = 1;
    }

    *R(VIRTIO_MMIO_QUEUE_READY) = 1;
    virtio_console_input_supply ();
}

void
virtio_console_input_supply () {
    int idx;
    char * buf;
    acquire (&cons.vcons_lock);
    do {
        idx = alloc_i_desc ();
    } while (idx == -1);
    buf = cons.inbuf[idx];
    memset (buf, 0, VIRTIO_CONS_BUF_SIZE);
    cons.i_desc[idx].addr = (uint64) buf;
    cons.i_desc[idx].flags = VRING_DESC_F_WRITE;
    cons.i_desc[idx].len = (uint32) VIRTIO_CONS_BUF_SIZE;
    cons.i_desc[idx].next = 0;
    cons.i_avail->ring[cons.i_avail->idx % NUM] = idx;
    cons.i_avail->idx += 1;
    // NOTE: NEED __sync_synchronize()
    cons.out_idx[idx] = 1;
    release (&cons.vcons_lock);
}

void
virtio_console_intr () {
    int id;
    *R(VIRTIO_MMIO_INTERRUPT_ACK) = *R(VIRTIO_MMIO_INTERRUPT_STATUS) & 0x3;
    while (cons.i_used_idx != cons.i_used->idx) {
        id = cons.i_used->ring[cons.i_used_idx % NUM].id;
        char *buf = (char *)cons.i_desc[id].addr;
        for (int i = 0; i < cons.i_used->ring[cons.i_used_idx % NUM].len; i++) {
            consoleintr (buf[i]);
        }
        free_i_desc (id);
        cons.i_used_idx += 1;
        virtio_console_input_supply ();
    }
    while (cons.o_used_idx != cons.o_used->idx) {
        id = cons.o_used->ring[cons.o_used_idx % NUM].id;
        free_o_desc (id);
        cons.out_idx[id] = 0; // done
        cons.o_used_idx += 1;
    }
}

void
virtio_console_output (char *buf, int len) {
    int idx;
    acquire (&cons.vcons_lock);
    do {
        idx = alloc_o_desc ();
    } while (idx == -1);
    cons.o_desc[idx].addr = (uint64) buf;
    cons.o_desc[idx].flags = 0;
    cons.o_desc[idx].len = (uint32) len;
    cons.o_desc[idx].next = 0;
    cons.o_avail->ring[cons.o_avail->idx % NUM] = idx;
    cons.o_avail->idx += 1;
    // NOTE: NEED __sync_synchronize()
    cons.out_idx[idx] = 1;
    *R(VIRTIO_MMIO_QUEUE_NOTIFY) = 1;
    /* waiting for interrupt */
    // while(cons.out_idx[idx]) {
    //     sleep (&cons.out_idx[idx], &cons.vcons_lock);
    // }
    // free_o_desc (idx);
    release (&cons.vcons_lock);
}
